import {Annotation, Command, END, interrupt, MemorySaver, START, StateGraph} from "@langchain/langgraph";


const StateAnnotation = Annotation.Root({
    input: Annotation<string>,
    userFeedback: Annotation<string>
})

const step1 = (_state: typeof StateAnnotation.State) => {
    console.log("---Step 1---")
    return {}
}

const humanFeedback = (_state: typeof StateAnnotation.State) => {
    console.log("--- humanFeedback ---")
    const feedback: string = interrupt("Please provide feedback")
    return {
        userFeedback: feedback
    }
}

const step3 = (_state: typeof StateAnnotation.State) => {
    console.log("--- Step 3---")
    return {}
}

const builder = new StateGraph(StateAnnotation)
.addNode("step1", step1)
.addNode('humanFeedback', humanFeedback)
.addNode('step3', step3)
.addEdge(START, 'step1')
.addEdge("step1",  'humanFeedback')
.addEdge('humanFeedback', 'step3')
.addEdge('step3', END)

const memory = new MemorySaver();

const graph = builder.compile({
    checkpointer: memory
})

const initialInput = {
    input: "hello world"
}

const config = {
    configurable: { thread_id: "1" }
}

for await (const event of await graph.stream(initialInput, config)) {
    console.log(event);
}

// Continue the graph execution
for await (const event of await graph.stream(
    new Command({ resume: "go to step 3! "}),
    config,
)) {
    console.log(event);
    console.log("\n====\n");
}